#!/usr/bin/env python3
"""
Knowledge Base Builder Script
知识库构建脚本

功能：构建COCO数据集的向量知识库和图像ID到描述的映射
"""

import os
import sys
import yaml
import logging
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from core.knowledge_base import KnowledgeBase


def setup_logging():
    """设置日志配置"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler('knowledge_base_build.log')
        ]
    )


def load_config(config_path: str) -> dict:
    """
    加载配置文件
    
    Args:
        config_path: 配置文件路径
        
    Returns:
        dict: 配置字典
    """
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        return config
    except Exception as e:
        logging.error(f"加载配置文件失败: {e}")
        raise


def main():
    """主函数"""
    try:
        # 设置日志
        setup_logging()
        logger = logging.getLogger(__name__)
        
        # 加载配置
        config_path = project_root / "configs" / "config.yaml"
        config = load_config(str(config_path))
        
        logger.info("开始构建COCO知识库...")
        
        # 创建知识库构建器
        kb_builder = KnowledgeBase(config)
        
        # 构建知识库
        kb_builder.build_knowledge_base(batch_size=32)
        
        logger.info("知识库构建完成！")
        
    except Exception as e:
        logging.error(f"知识库构建失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
